package websocket

import (
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
)

// 全局连接管理
var upgrader = websocket.Upgrader{}
var connections = make(map[string]*websocket.Conn)
var mu sync.Mutex

// HandleWebSocket 处理 WebSocket 连接
func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 升级 WebSocket 连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		http.Error(w, "Could not upgrade connection", http.StatusBadRequest)
		return
	}
	// 从查询参数获取 token
	token := r.URL.Query().Get("token")
	mu.Lock()
	connections[token] = conn // 保存连接
	mu.Unlock()

	// 处理来自客户端的消息
	defer func() {
		mu.Lock()
		delete(connections, token)
		mu.Unlock()
		conn.Close()
	}()

	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

// SendMessage 向特定 token 的客户端发送消息
func SendMessage(token string, message string) {
	mu.Lock()
	conn, ok := connections[token]
	mu.Unlock()

	if ok {
		conn.WriteMessage(websocket.TextMessage, []byte(message))
	}
}
